#!/usr/bin/python

"""
RSA 加密/解密, 签名/验证 接口
"""

from binascii import unhexlify
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5
from Crypto.Signature import PKCS1_v1_5 as SignPKCS1_v1_5
import base64
from Crypto.Hash import SHA256


def create_rsa_key(password: str = None,
                   bits: int = 1024,
                   key_format: str = 'PEM',
                   pkcs: (1, 8) = 1) -> (bytes, bytes):
    """
    创建rsa公私钥
    @params password: 密码
    @params bits: 密钥长度
    @params key_format: 密钥格式
    @params pkcs:
    @return : 私钥，公钥
    """
    key = RSA.generate(bits)
    private_key = key.exportKey(
        format=key_format, passphrase=password, pkcs=pkcs)
    public_key = key.publickey().exportKey()
    return private_key, public_key


def encrypt(data: bytes, public_key: str):
    """
    rsa加密数据
    @params data: 待加密数据
    @params public_key: 公钥
    @return : 密文的base64
    """
    recipient_key = RSA.importKey(public_key)
    cipher_rsa = PKCS1_v1_5.new(recipient_key)
    return base64.b64encode(cipher_rsa.encrypt(data))


def decrypt(data: str, private_key: str, password: str = None) -> bytes:
    """
    rsa解密数据
    @params data: 密文的base64
    @params private_key: 秘钥
    @params password: 密码
    @return : 明文
    """
    priv_key = RSA.importKey(private_key, passphrase=password)
    cipher_rsa = PKCS1_v1_5.new(priv_key)
    return cipher_rsa.decrypt(base64.b64decode(data), None)


def rsa_sign(data: bytes, private_key: str, password: str = None):
    """
    rsa签名
    @params data: 待签名数据
    @password: 签名密码
    @return : 签名
    """
    priv_key = RSA.importKey(private_key, passphrase=password)
    hash_obj = SHA256.new(data)
    signature = base64.b64encode(SignPKCS1_v1_5.new(priv_key).sign(hash_obj))
    return signature


def rsa_verify(data, public_key, signature) -> bool:
    """
    rsa验证签名
    @params data: 签名内容
    @params public_key: 公钥
    @params signature: 签名
    @return : 签名是否合法 True,合法   False非法
    """
    pub_key = RSA.importKey(public_key)
    hash_obj = SHA256.new(data)
    try:
        # 因为签名被base64编码，所以这里先解码，再验签
        SignPKCS1_v1_5.new(pub_key).verify(
            hash_obj, base64.b64decode(signature))
        return True
    except (ValueError, TypeError):
        return False


if __name__ == '__main__':
    # 创建私钥以及公钥
    priv_key, pub_key = create_rsa_key()
    priv_key = priv_key.decode('utf8')
    pub_key = pub_key.decode('utf8')
    # 加密使用公钥
    enc_data = encrypt(b'qianyu', pub_key)

    # 解密使用私钥
    print(decrypt(enc_data, priv_key))

    # 生成签名使用私钥
    sign = rsa_sign(b'123456abd', priv_key)
    # 检验签名使用公钥
    print(rsa_verify(b'123456abd', pub_key, sign))
